#include "core/memory.h"
#include "ipc/mutex.h"
#include "tools/log.h"
#include "tools/klib.h"
#include "cpu/mmu.h"

static addr_alloc_t paddr_alloc;

static PDE_T kernel_page_dir[PDE_CNT] __attribute__((aligned(MEM_PAGE_SIZE)));

static void addr_alloc_init(addr_alloc_t *alloc, 
                            uint8_t *bits, 
                            uint32_t start, 
                            uint32_t size, 
                            uint32_t page_size)
{
    mutex_init(&alloc->mutex);
    alloc->start = start;
    alloc->size = size;
    alloc->page_size = page_size;
    bitmap_init(&alloc->bitmap, bits, alloc->size / page_size, 0);
}

static uint32_t addr_alloc_page(addr_alloc_t *alloc, int page_count)
{
    // 分配位图需要加锁
    mutex_lock(&alloc->mutex);
    int page_index = bitmap_alloc_nbits(&alloc->bitmap, 0, page_count);
    mutex_unlock(&alloc->mutex);

    uint32_t addr = 0;
    if (page_index >= 0) {
        addr = alloc->start + page_index * alloc->page_size;
    }
    return addr;
}

static void addr_free_page(addr_alloc_t *alloc, uint32_t addr, int page_count)
{
    uint32_t pg_idx = (addr - alloc->start) / alloc->page_size;

    mutex_lock(&alloc->mutex);
    bitmap_set_bit(&alloc->bitmap, pg_idx, page_count, 0);
    mutex_unlock(&alloc->mutex);
}

void show_memory_info(boot_info_t *boot_info)
{
    log_printf("mem region:");
    for (int i = 0; i < boot_info->ram_region_count; i++) {
        log_printf("[%d]: 0x%x - 0x%x", i, 
            boot_info->ram_region_cfg[i].start, 
            boot_info->ram_region_cfg[i].size);
    }
    log_printf("\n");
}

static uint32_t total_mem_size(boot_info_t *boot_info)
{
    uint32_t mem_size = 0;
    for (int i = 0; i < boot_info->ram_region_count; i++) {
        mem_size += boot_info->ram_region_cfg[i].size;
    }
    return mem_size;
}

PTE_T * find_pte(PDE_T *page_dir, uint32_t vaddr, int alloc)
{
    PTE_T *page_table;
    //页目录表 + 虚拟地址对应PDT中的索引 = 页目录项地址
    PDE_T *pde = page_dir + pde_index(vaddr);
    if (pde->P) {   //如果页目录项存在，则取出其物理地址
        page_table = (PTE_T *)pde_paddr(pde);
    } else {
        //如果页目录项不存在且无需分配，直接返回
        if (alloc == 0) {
            return NULL;
        }

        //否则分配物理页
        uint32_t pg_paddr = addr_alloc_page(&paddr_alloc, 1);
        if (pg_paddr == 0) {
            return NULL;
        }

        pde->v = pg_paddr | PDE_P | PDE_W | PDE_U;
        page_table = (PTE_T *)pg_paddr;
        kernel_memset(page_table, 0, MEM_PAGE_SIZE);
    }

    return page_table + pte_index(vaddr);
}

int memory_create_map(PDE_T *page_dir, uint32_t vaddr, uint32_t paddr, 
                            int page_count, uint32_t perm)
{
    for (int i = 0; i < page_count; i++) {
        // log_printf("create map: v-0x%x, p-0x%x, perm: 0x%x", vaddr, paddr, perm);

        PTE_T *pte = find_pte(page_dir, vaddr, 1);
        if (pte == NULL) {
            log_printf("create pte failed. pte == 0");
            return -1;
        }
        
        // log_printf("pte addr:0x%x", (uint32_t)pte);

        ASSERT(pte->P == 0);
        pte->v = paddr | perm | PTE_P;

        vaddr += MEM_PAGE_SIZE;
        paddr += MEM_PAGE_SIZE;
    }
    return 0;
}

void create_kernel_table()
{
    extern uint8_t s_text[];
    extern uint8_t e_text[];
    extern uint8_t s_data[];
    extern uint8_t kernel_base[];

    static memory_map_t kernel_map[] = {
        {kernel_base, s_text, kernel_base, PTE_W},
        {s_text, e_text, s_text, 0},
        {s_data, (void *)MEM_EBDA_START, s_data, PTE_W},
        {(void *)MEM_EXT_START, (void *)MEM_EXT_END, (void *)MEM_EXT_START, PTE_W}
    };

    for (int i = 0; i < sizeof(kernel_map) / sizeof(memory_map_t); i++) {
        uint32_t vstart = down2((uint32_t)kernel_map[i].vstart, MEM_PAGE_SIZE);
        uint32_t vend = up2((uint32_t)kernel_map[i].vend, MEM_PAGE_SIZE);
        uint32_t paddr = down2((uint32_t)(kernel_map[i].pstart), MEM_PAGE_SIZE);
        int page_count = (vend - vstart) / MEM_PAGE_SIZE;

        memory_create_map(kernel_page_dir, vstart, paddr, page_count, kernel_map[i].perm);
    }

}

void memory_init(boot_info_t *boot_info)
{
    // kernel.lds中定义的符号，在.bss段之后
    extern uint8_t *mem_free_start;  

    log_printf("mem init");
    show_memory_info(boot_info);

    // 1MB以内空闲内存起始地址
    uint8_t *p_mem_free = (uint8_t *)&mem_free_start;    

    // 计算1MB以上空闲内存空间，并按页大小对齐
    uint32_t mem_up1MB_free = total_mem_size(boot_info) - MEM_EXT_START;
    mem_up1MB_free = down2(mem_up1MB_free, MEM_PAGE_SIZE);

    log_printf("free memory: 0x%x, size: 0x%x", MEM_EXT_START, mem_up1MB_free);
    // 初始化地址分配器，设置在.bss段之后放置bitmap
    addr_alloc_init(&paddr_alloc, p_mem_free, MEM_EXT_START, mem_up1MB_free, MEM_PAGE_SIZE);
    // 计算1MB以上空闲内存所需的bitmap字节数，将p_mem_free指向bitmap之后
    p_mem_free += bitmap_byte_count(mem_up1MB_free / MEM_PAGE_SIZE);

    ASSERT(p_mem_free < (uint8_t *)MEM_EBDA_START);

    create_kernel_table();
    mmu_set_page_dir((uint32_t)kernel_page_dir);
}

uint32_t memory_create_uvm()
{
    PDE_T *page_dir = (PDE_T *)addr_alloc_page(&paddr_alloc, 1);
    if (page_dir == NULL) {
        return 0;
    }

    kernel_memset((void *)page_dir, 0, MEM_PAGE_SIZE);
    uint32_t user_pde_start = pde_index(MEMORY_TASK_BASE);
    for (int i = 0; i < user_pde_start; i++) {
        page_dir[i].v = kernel_page_dir[i].v;
    }

    return (uint32_t)page_dir;
}

int memory_alloc_page_for_pagedir(uint32_t pagedir, uint32_t addr, 
                                    uint32_t size, int perm)
{
    uint32_t curr_vaddr = addr;
    int page_count = up2(size, MEM_PAGE_SIZE) / MEM_PAGE_SIZE;
    for (int i = 0; i < page_count; i++) {
        uint32_t paddr = addr_alloc_page(&paddr_alloc, 1);
        if (paddr == 0) {
            log_printf("Memory alloc failed. No memory.");
            return 0;
        }

        int err = memory_create_map((PDE_T *)pagedir, curr_vaddr, paddr, 
                                    1, perm);
        if (err < 0) {
            log_printf("Create memory failed. err = %d", err);
            return 0;
        }
        curr_vaddr += MEM_PAGE_SIZE;
    }   
    return 0;
}

int memory_alloc_page_for(uint32_t addr, uint32_t size, int perm)
{
    return memory_alloc_page_for_pagedir(get_current_task()->tss.cr3, 
                                            addr, size, perm);
}